

use std::{process::Command, fs, os::unix::prelude::PermissionsExt};

pub trait IpBlock {

    /**
     * 初始化
     */
    fn init(ipset_path: &str) -> Self;

    /**
     * 添加阻断IP
     */
    fn add(&self, ip: &str) -> bool;

    /**
     * 添加阻断IP，设置超时时间
     */
    fn add_timeout(&self, ip: &str, timeout: i32) -> bool;

    /**
     * 删除阻断IP
     */
    fn del(&self, ip: &str) -> bool;

    /**
     * 清空阻断IP
     */
    fn flush(&self) -> bool;

    /**
     * 列出所有阻断IP
     */
    fn list(&self) -> Option<String>;

    /**
     * 检查依赖
     */
    fn check_dependencies(&self) -> bool;

    /**
     * 检查iptables
     */
    fn check_iptables(&self) -> bool;

    /**
     * 检查ipset
     */
    fn check_ipset(&self) -> bool;

}

#[derive(Clone)]
pub struct Iptables {

    /**
     * "TAILMON-EDR-BLOCK-IP"
     */
    ipset_name: String,

    /**
     * "hash:net"
     */
    ipset_type: String,

    /**
     * "ipset"
     */
    ipset_path: String,
}

impl IpBlock for Iptables {

    /**
     * 初始化
     */
    fn init(ipset_path: &str) -> Self {
        let plugin_root = std::fs::canonicalize(ipset_path).unwrap();
        let plugin_root = plugin_root.to_str().unwrap();

        let iptables = Iptables {
            ipset_name: "TAILMON-EDR-BLOCK-IP".to_string(),
            ipset_type: "hash:net".to_string(),
            ipset_path: plugin_root.to_string(),
        };
        if !iptables.check_dependencies() {
            panic!("Iptables check_dependencies failed")
        }
        iptables.check_ipset();
        iptables.check_iptables();
        iptables
    }

    /**
     * 添加阻断IP
     * ipset add TAILMON-EDR-BLOCK-IP 192.168.1.1
     */
    fn add(&self, ip: &str) -> bool {
        let status = Command::new(format!("{}/bin/ipset", &self.ipset_path))
                            .current_dir(&self.ipset_path)
                            .env("LD_LIBRARY_PATH", format!("{}/lib", &self.ipset_path))
                            .arg("add")
                            .arg(&self.ipset_name)
                            .arg(ip)
                            .status()
                            .expect("failed to execute process");

        println!("process finished with: {}", status);

        status.success()
    }

    /**
     * 添加阻断IP，设置超时时间
     * ipset add TAILMON-EDR-BLOCK-IP 192.168.1.1 timeout 1000
     */
    fn add_timeout(&self, ip: &str, timeout: i32) -> bool {
        
        let status = Command::new(format!("{}/bin/ipset", &self.ipset_path))
                            .current_dir(&self.ipset_path)
                            .env("LD_LIBRARY_PATH", format!("{}/lib", &self.ipset_path))
                            .arg("add")
                            .arg(&self.ipset_name)
                            .arg(ip)
                            .arg("timeout")
                            .arg(timeout.to_string())
                            .status()
                            .expect("failed to execute process");

        println!("process finished with: {}", status);

        status.success()
    }

    /**
     * 删除阻断IP
     * ipset del TAILMON-EDR-BLOCK-IP 192.168.1.1
     */
    fn del(&self, ip: &str) -> bool {
        let status = Command::new(format!("{}/bin/ipset", &self.ipset_path))
                .current_dir(&self.ipset_path)
                .env("LD_LIBRARY_PATH", format!("{}/lib", &self.ipset_path))
                .arg("del")
                .arg(&self.ipset_name)
                .arg(ip)
                .status()
                .expect("failed to execute process");

        println!("process finished with: {}", status);

        status.success()
    }

    /**
     * 清空阻断IP
     * ipset flush TAILMON-EDR-BLOCK-IP
     */
    fn flush(&self) -> bool {
        let status = Command::new(format!("{}/bin/ipset", &self.ipset_path))
                .current_dir(&self.ipset_path)
                .env("LD_LIBRARY_PATH", format!("{}/lib", &self.ipset_path))
                .arg("flush")
                .arg(&self.ipset_name)
                .status()
                .expect("failed to execute process");

        println!("process finished with: {}", status);

        status.success()
    }

    /**
     * 列出所有阻断IP
     * ipset list TAILMON-EDR-BLOCK-IP
     */
    fn list(&self) -> Option<String> {
        let output = Command::new(format!("{}/bin/ipset", &self.ipset_path))
            .current_dir(&self.ipset_path)
            .env("LD_LIBRARY_PATH", format!("{}/lib", &self.ipset_path))
            .arg("list")
            .arg(&self.ipset_name)
            .output()
            .expect("Failed to execute command");

        let out = String::from_utf8_lossy(output.stdout.as_slice()).to_string();
        Some(out)
    }

    /**
     * 检查依赖
     * 
     */
    fn check_dependencies(&self) -> bool {
        let path = format!("{}/bin/ipset", &self.ipset_path);
        let ipset_bin = std::path::Path::new(&path);
        let path = format!("{}/lib/libipset.so.2", &self.ipset_path);
        let libipset_so = std::path::Path::new(&path);
        let path = format!("{}/lib/libmnl.so.0", &self.ipset_path);
        let libmnl_so = std::path::Path::new(&path);
        if !ipset_bin.exists() || !libipset_so.exists() || !libmnl_so.exists() {
            return false;
        } else {
            fs::set_permissions(ipset_bin, fs::Permissions::from_mode(0o755)).unwrap();
            return true;
        }
    }

    /**
     * 检查iptables
     */
    fn check_iptables(&self) -> bool {
        let output = Command::new("iptables")
            .arg("-L")
            .arg("INPUT")
            .arg("--line-numbers")
            .output()
            .expect("Failed to execute command");

        let out = String::from_utf8_lossy(output.stdout.as_slice()).to_string();

        let lines = out.split("\n");
        for line in lines {
            if line.contains(&self.ipset_name) {
                let seperator = regex::Regex::new(r"\s+").expect("Invalid regex");
                let splits: Vec<_> = seperator.split(line).into_iter().collect();
                let order = splits[0];
                if order != "1" {
                    let status = Command::new("iptables")
                        .arg("-D")
                        .arg("INPUT")
                        .arg(order)
                        .status()
                        .expect("failed to execute process");

                    println!("process finished with: {}", status);
                } else {
                    return true;
                }
            }
        }

        let status = Command::new("iptables")
            .arg("-I")
            .arg("INPUT")
            .arg("-m")
            .arg("set")
            .arg("--match-set")
            .arg(&self.ipset_name)
            .arg("src")
            .arg("-j")
            .arg("DROP")
            .status()
            .expect("failed to execute process");

        println!("process finished with: {}", status);
        return status.success();
    }

    /**
     * 检查ipset
     */
    fn check_ipset(&self) -> bool {
        let status = Command::new(format!("{}/bin/ipset", &self.ipset_path))
            .current_dir(&self.ipset_path)
            .env("LD_LIBRARY_PATH", format!("{}/lib", &self.ipset_path))
            .arg("create")
            .arg(&self.ipset_name)
            .arg(&self.ipset_type)
            .arg("timeout")
            .arg("0")
            .status()
            .expect("failed to execute process");

        println!("process finished with: {}", status);
        return status.success();
    }

}